import os
import cv2
import random


def split(fname_im, fname_gt, tr_outpth_im, val_outpth_im, tr_outpth_gt, val_outpth_gt, tr_idx, val_idx):
    img_lst = os.listdir(fname_im)
    gt_lst = os.listdir(fname_gt)
    img_lst.sort()
    gt_lst.sort()
    cnt = 0

    for tr_id in tr_idx:
        img_name = img_lst[tr_id]
        img = cv2.imread(fname_im + '/' + img_name)
        cv2.imwrite(tr_outpth_im + '/' + img_name, img)

        gt_name = gt_lst[tr_id]
        with open(fname_gt + '/' + gt_name, 'r', encoding='utf-8') as f:
            txt = f.readlines()
        with open(tr_outpth_gt + '/gt_' + gt_name, 'w', encoding='utf-8') as f:
            f.writelines(txt)

        cnt += 1
        if cnt % 100 == 0:
            print(f"finished:{cnt}")

    for val_id in val_idx:
        # print(val_id)
        img_name = img_lst[val_id]
        img = cv2.imread(fname_im + '/' + img_name)
        cv2.imwrite(val_outpth_im + '/' + img_name, img)

        gt_name = gt_lst[val_id]
        print(img_name, gt_name)
        with open(fname_gt + '/' + gt_name, 'r', encoding='utf-8') as f:
            txt = f.readlines()
        with open(val_outpth_gt + '/gt_' + gt_name, 'w', encoding='utf-8') as f:
            f.writelines(txt)

        cnt += 1
        if cnt % 100 == 0:
            print(f"finished:{cnt}")
    print("img end")

# def split_gt(fname, tr_outpth, val_outpth, p):
#     print("gt begin")
#     random.seed(1)
#     gt_lst = os.listdir(fname)
#     split_point = int(len(gt_lst) * p)
#
#     random.shuffle(gt_lst)
#     tr_lst = gt_lst[:split_point]
#     val_lst = gt_lst[split_point:]
#     cnt = 0
#     for gt_name in tr_lst:
#         with open(fname + '/' + gt_name, 'r', encoding='utf-8') as f:
#             txt = f.readlines()
#         with open(tr_outpth + '/gt_' + gt_name, 'w', encoding='utf-8') as f:
#             f.writelines(txt)
#         cnt += 1
#         if cnt % 100 == 0:
#              print(f"finished:{cnt}")
#
#     for gt_name in val_lst:
#         with open(fname + '/' + gt_name, 'r', encoding='utf-8') as f:
#             txt = f.readlines()
#         with open(val_outpth + '/gt_' + gt_name, 'w', encoding='utf-8') as f:
#             f.writelines(txt)
#         cnt = cnt + 1
#         if cnt % 100 == 0:
#             print(f"finished:{cnt}")
#     print("gt end")
if __name__ == '__main__':
    org_im = 'img'
    org_gt = 'gt'
    img_tr = 'img_tr'
    img_val = 'img_val'
    gt_tr = 'gt_tr' 
    gt_val = 'gt_val'

    num = len(os.listdir(org_im))
    random.seed(1)
    idx = list(range(num))
    random.shuffle(idx)
    split_point = int(num * 0.7)
    tr_idx = idx[:split_point]
    val_idx = idx[split_point:]

    if not os.path.exists('img_tr'):
        os.mkdir(img_tr)
        os.mkdir(img_val)
        os.mkdir(gt_tr)
        os.mkdir(gt_val)
    split(org_im, org_gt, img_tr, img_val, gt_tr, gt_val, tr_idx, val_idx)
